from .image_convert_base import ConvertBase
import cv2 as cv
import random
import numpy as np

class RotationConvert(ConvertBase):
    
    #is_center_rotation ： true 是否是中心点旋转
    # 如果设置了angle_list 就优先从angle_list 种随机一个角度来旋转
    def __init__(self, center = None, use_rate = 0.5, min_angle = 0, max_angle = 360, angle_list = None, is_rand = True, angle = 0, is_center_rotation = True):
        super().__init__(use_rate)
        
        self.angle = angle
        self.center = center
        self.is_rand = is_rand
        self.min_angle = min_angle
        self.max_angle = max_angle
        
        self.is_center_rotation = is_center_rotation
        if angle_list is not None and isinstance(angle_list, list) == False:
            raise Exception('angle_list 必须是一个角度的列表')
        self.angle_list = angle_list
    
    def __get_angle(self):
        if self.angle_list is not None:
            return random.sample(self.angle_list, 1)[0]
        
        return np.random.random() * (self.max_angle - self.min_angle) + self.min_angle
        
        
    def convert(self, img, bboxes, points):
        if self.is_rand:
            self.angle = self.__get_angle()

        h, w, _ = img.shape
        size = (w, h)
        if self.is_center_rotation:
            center = (w / 2, h / 2)
            
        #进行旋转
        rotate = cv.getRotationMatrix2D(center, self.angle, 1)
        img = cv.warpAffine(img, rotate, size)
        if bboxes is not None:
            # #旋转区域框
            bboxes = self.rotation_rect_batch(bboxes, rotate, size)
        if points is not None:
            #旋转关键点
            points = self.rotation_points(points, rotate, size)
        
        return img, bboxes, points
    
        #旋转矩阵
    def rotation_rect_batch(self, rect, rotation, size):
        # print('\nrect_shape : ', rect.shape)
        #转换成点的形式
        rect = rect.reshape([-1, 2, 2])
        #转换成4个点
        r = np.ones([rect.shape[0], 4, 2])
        r[:, 0:2] = rect[:]
        # print('r : \n', r[0:2])
        r[:, 2] = r[:, 0]
        r[:, 2, 0] = r[:, 1, 0]
        r[:, 3] = r[:, 1]
        r[:, 3, 0] = r[:, 0, 0]
        #拼接
        r = r.reshape([-1, 2])
        # res = np.matmul(rotation, r)
        res = self.rotation_points(r, rotation, size)
        res = res.reshape([rect.shape[0], -1, 2])
        # print('shape : ', res.shape, 'res : \n', res)
        # print(np.max(res[0]))
        max_value = np.max(res, axis = 1)
        min_value = np.min(res, axis = 1)
        #
        res = np.hstack((min_value, max_value))
        res = np.maximum(res, 0)
        res[:, [0, 2]] = np.minimum(res[:, [0, 2]], size[0])
        res[:, [1, 3]] = np.minimum(res[:, [1, 3]], size[1])
        return res
    
    #旋转所有的点坐在的坐标
    def rotation_points(self, points, rotation, size):
        w, h = size
        #进行转置矩阵 
        r = np.transpose(points)
        target = np.ones([r.shape[0] + 1, r.shape[1]], r.dtype)
        # #用于插入
        # data = np.ones([1, r.shape[1]])
        # #新增一行
        # r = np.insert(r, r.shape[0], values = data, axis = 0)
        target[0:r.shape[0], :] = r
        r = target
        #进行旋转
        res = np.matmul(rotation, r)
        res[:, 0] = np.minimum(res[:, 0], w)
        res[:, 0] = np.maximum(res[:, 0], 0)
        res[:, 1] = np.minimum(res[:, 1], h)
        res[:, 1] = np.maximum(res[:, 1], 0)
        
        #转置回去
        res =  np.transpose(res)
        return res
    
    
    # 坐标旋转
    def rotation_boxes(self, boxes, marix, size):
        convert_boxes = np.ones([boxes.shape[0], 3, 4], np.float32)
        convert_boxes[:, [0, 1], 0] = boxes[:, [0, 1]]
        convert_boxes[:, [0, 1], 1] = boxes[:, [2, 1]]
        convert_boxes[:, [0, 1], 2] = boxes[:, [0, 3]]
        convert_boxes[:, [0, 1], 3] = boxes[:, [2, 3]]
        
        res = np.matmul(marix, convert_boxes)
        res = np.transpose(res).reshape([-1, 4, 2])
        
        max_value = np.max(res, axis = 1)
        min_value = np.min(res, axis = 1)
        #
        res = np.hstack((min_value, max_value))
       
        res = np.maximum(res, 0)
        res[:, [0, 2]] = np.minimum(res[:, [0, 2]], size[0])
        res[:, [1, 3]] = np.minimum(res[:, [1, 3]], size[1])
        
        return res
    
    